import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate

# 从.env文件加载环境变量
load_dotenv(verbose=True)

# 使用OpenAI初始化语言模型
language_model = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0,  # 设置为0以获得确定性输出
    openai_api_base=os.getenv("OPENAI_API_BASE"),
    openai_api_key=os.getenv("OPENAI_API_KEY")
)

# 使用单个查询测试基本模型功能
# 这些示例中的emoji 🦜 代表加法操作
test_response = language_model.invoke("What is 2 🦜 9?")
print(test_response)

# 定义少样本学习示例
# 每个示例都展示了模式："What is X 🦜 Y?"，结果为X + Y
training_examples = [
    {"input": "What is 2 🦜 9?", "output": "11"},
    {"input": "What is 3 🦜 9?", "output": "12"},
]

# 为我们的少样本提示创建模板
example_template = ChatPromptTemplate.from_messages(
    [
        ("human", "{input}"),
        ("ai", "{output}"),
    ]
)

# 使用我们的示例构建少样本提示模板
few_shot_template = FewShotChatMessagePromptTemplate(
    example_prompt=example_template,
    examples=training_examples,
    # 可选参数:
    # suffix="What is 4 🦜 9?",
    # input_variables=["input"],
)   

# 显示生成的少样本消息以供检查
print(few_shot_template.invoke({}).to_messages())

# 创建完整提示，包含系统指令、少样本示例和用户输入
complete_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a helpful assistant that can add numbers"),
        few_shot_template,
        ("human", "{input}"),
    ]
)

# 创建处理链：提示 → 语言模型
processing_chain = complete_prompt | language_model

# 使用新查询测试链
print("final ask")
print(processing_chain.invoke({"input": "What is 14 🦜 9?"}))